
namespace Microsoft.VsSDK.IntegrationTestLibrary
{
   using System;
   using System.Collections.Generic;
   using System.Text;
   using System.Runtime.InteropServices;
   using System.Threading;
   using Microsoft.VisualStudio.Shell.Interop;
   using Microsoft.VisualStudio.Shell;

   /// <summary>
   /// This class is responsible to close dialog boxes that pop up during different VS Calls
   /// </summary>
   internal class DialogBoxPurger : IDisposable
   {
      /// <summary>
      /// The default number of milliseconds to wait for the threads to signal to terminate.
      /// </summary>
      private const int DefaultMillisecondsToWait = 3500;

      /// <summary>
      /// Object used for synchronization between thread calls.
      /// </summary>
      internal static volatile object Mutex = new object();

      /// <summary>
      /// The IVsUIShell. This cannot be queried on the working thread from the service provider. Must be done in the main thread.!!
      /// </summary>
      private IVsUIShell uiShell;

      /// <summary>
      /// The button to "press" on the dialog.
      /// </summary>
      private int buttonAction;

      /// <summary>
      /// Thread signales to the calling thread that it is done.
      /// </summary>
      private bool exitThread = false;

      /// <summary>
      /// Calling thread signales to this thread to die.
      /// </summary>
      private AutoResetEvent threadDone = new AutoResetEvent(false);

      /// <summary>
      /// The queued thread started.
      /// </summary>
      private AutoResetEvent threadStarted = new AutoResetEvent(false);

      /// <summary>
      /// The result of the dialogbox closing for all the dialog boxes. That is if there are two of them and one fails this will be false.
      /// </summary>
      private bool dialogBoxCloseResult = false;

      /// <summary>
      /// The expected text to see on the dialog box. If set the thread will continue finding the dialog box with this text.
      /// </summary>
      private string expectedDialogBoxText = String.Empty;

      /// <summary>
      /// The number of the same  dialog boxes to wait for.
      /// This is for scenarios when two dialog boxes with the same text are popping up.
      /// </summary>
      private int numberOfDialogsToWaitFor = 1;

      /// <summary>
      /// Has the object been disposed.
      /// </summary>
      private bool isDisposed;

      /// <summary>
      /// Overloaded ctor.
      /// </summary>
      /// <param name="buttonAction">The botton to "press" on the dialog box.</param>
      /// <param name="numberOfDialogsToWaitFor">The number of dialog boxes with the same message to wait for. This is the situation when the same action pops up two of the same dialog boxes</param>
      /// <param name="expectedDialogMesssage">The expected dialog box message to check for.</param>
      internal DialogBoxPurger(int buttonAction, int numberOfDialogsToWaitFor, string expectedDialogMesssage)
      {
         this.buttonAction = buttonAction;
         this.numberOfDialogsToWaitFor = numberOfDialogsToWaitFor;
         this.expectedDialogBoxText = expectedDialogMesssage;
      }

      /// <summary>
      /// Overloaded ctor.
      /// </summary>
      /// <param name="buttonAction">The botton to "press" on the dialog box.</param>
      /// <param name="numberOfDialogsToWaitFor">The number of dialog boxes with the same message to wait for. This is the situation when the same action pops up two of the same dialog boxes</param>
      internal DialogBoxPurger(int buttonAction, int numberOfDialogsToWaitFor)
      {
         this.buttonAction = buttonAction;
         this.numberOfDialogsToWaitFor = numberOfDialogsToWaitFor;
      }

      /// <summary>
      /// Overloaded ctor.
      /// </summary>
      /// <param name="buttonAction">The botton to "press" on the dialog box.</param>
      /// <param name="expectedDialogMesssage">The expected dialog box message to check for.</param>
      internal DialogBoxPurger(int buttonAction, string expectedDialogMesssage)
      {
         this.buttonAction = buttonAction;
         this.expectedDialogBoxText = expectedDialogMesssage;
      }

      /// <summary>
      /// Overloaded ctor.
      /// </summary>
      /// <param name="buttonAction">The botton to "press" on the dialog box.</param>
      internal DialogBoxPurger(int buttonAction)
      {
         this.buttonAction = buttonAction;
      }

      /// <summary>
      #region IDisposable Members

      void IDisposable.Dispose()
      {
         if (this.isDisposed)
         {
            return;
         }

         this.WaitForDialogThreadToTerminate();

         this.isDisposed = true;
      }

      /// <summary>
      /// Spawns a thread that will start listening to dialog boxes.
      /// </summary>
      internal void Start()
      {
         // We ask for the uishell here since we cannot do that on the therad that we will spawn.
         IVsUIShell uiShell = Package.GetGlobalService(typeof(SVsUIShell)) as IVsUIShell;

         if (uiShell == null)
         {
            throw new InvalidOperationException("Could not get the uiShell from the serviceProvider");
         }

         this.uiShell = uiShell;

         System.Threading.Thread thread = new System.Threading.Thread(new ThreadStart(this.HandleDialogBoxes));
         thread.Start();

         // We should never deadlock here, hence do not use the lock. Wait to be sure that the thread started.
         this.threadStarted.WaitOne(3500, false);
      }

      /// <summary>
      /// Waits for the dialog box close thread to terminate. If the thread does not signal back within millisecondsToWait that it is shutting down,
      /// then it will tell to the thread to do it.
      /// </summary>
      internal bool WaitForDialogThreadToTerminate()
      {
         return this.WaitForDialogThreadToTerminate(DefaultMillisecondsToWait);
      }

      /// <summary>
      /// Waits for the dialog box close thread to terminate. If the thread does not signal back within millisecondsToWait that it is shutting down,
      /// then it will tell to the thread to do it.
      /// </summary>
      /// <param name="millisecondsToWait">The number milliseconds to wait for until the dialog purger thread is signaled to terminate. This is just for safe precaution that we do not hang. </param>
      /// <returns>The result of the dialog boxes closing</returns>
      internal bool WaitForDialogThreadToTerminate(int numberOfMillisecondsToWait)
      {
         bool signaled = false;

         // We give millisecondsToWait sec to bring up and close the dialog box.
         signaled = this.threadDone.WaitOne(numberOfMillisecondsToWait, false);

         // Kill the thread since a timeout occured.
         if (!signaled)
         {
            lock (Mutex)
            {
               // Set the exit thread to true. Next time the thread will kill itselfes if it sees 
               this.exitThread = true;
            }

            // Wait for the thread to finish. We should never deadlock here.
            this.threadDone.WaitOne();
         }

         return this.dialogBoxCloseResult;
      }

      /// <summary>
      /// This is the thread method. 
      /// </summary>
      private void HandleDialogBoxes()
      {
         // No synchronization numberOfDialogsToWaitFor since it is readonly
         IntPtr[] hwnds = new IntPtr[this.numberOfDialogsToWaitFor];
         bool[] dialogBoxCloseResults = new bool[this.numberOfDialogsToWaitFor];

         try
         {
            // Signal that we started
            lock (Mutex)
            {
               this.threadStarted.Set();
            }

            // The loop will be exited either if a message is send by the caller thread or if we found the dialog. If a message box text is specified the loop will not exit until the dialog is found.
            bool stayInLoop = true;
            int dialogBoxesToWaitFor = 1;

            while (stayInLoop)
            {
               int hwndIndex = dialogBoxesToWaitFor - 1;

               // We need to lock since the caller might set context to null.
               lock (Mutex)
               {
                  if (this.exitThread)
                  {
                     break;
                  }

                  // We protect the shell too from reentrency.
                  this.uiShell.GetDialogOwnerHwnd(out hwnds[hwndIndex]);

               }

               if (hwnds[hwndIndex] != IntPtr.Zero)
               {
                  StringBuilder windowClassName = new StringBuilder(256);
                  NativeMethods.GetClassName(hwnds[hwndIndex], windowClassName, windowClassName.Capacity);

                  // The #32770 is the class name of a messagebox dialog.
                  if (windowClassName.ToString().Contains("#32770"))
                  {
                     IntPtr unmanagedMemoryLocation = IntPtr.Zero;
                     string dialogBoxText = String.Empty;
                     try
                     {
                        unmanagedMemoryLocation = Marshal.AllocHGlobal(10 * 1024);
                        NativeMethods.EnumChildWindows(hwnds[hwndIndex], new NativeMethods.CallBack(FindMessageBoxString), unmanagedMemoryLocation);
                        dialogBoxText = Marshal.PtrToStringUni(unmanagedMemoryLocation);
                     }
                     finally
                     {
                        if (unmanagedMemoryLocation != IntPtr.Zero)
                        {
                           Marshal.FreeHGlobal(unmanagedMemoryLocation);
                        }
                     }

                     lock (Mutex)
                     {

                        // Since this is running on the main thread be sure that we close the dialog.
                        bool dialogCloseResult = false;
                        if (this.buttonAction != 0)
                        {
                           dialogCloseResult = NativeMethods.EndDialog(hwnds[hwndIndex], this.buttonAction);
                        }

                        // Check if we have found the right dialog box.
                        if (String.IsNullOrEmpty(this.expectedDialogBoxText) || (!String.IsNullOrEmpty(dialogBoxText) && String.Compare(this.expectedDialogBoxText, dialogBoxText.Trim(), StringComparison.OrdinalIgnoreCase) == 0))
                        {
                           dialogBoxCloseResults[hwndIndex] = dialogCloseResult;
                           if (dialogBoxesToWaitFor++ >= this.numberOfDialogsToWaitFor)
                           {
                              stayInLoop = false;
                           }
                        }
                     }
                  }
               }
            }
         }
         finally
         {
            //Let the main thread run a possible close command.
            System.Threading.Thread.Sleep(2000);

            foreach (IntPtr hwnd in hwnds)
            {
               // At this point the dialog should be closed, if not attempt to close it.
               if (hwnd != IntPtr.Zero)
               {
                  NativeMethods.SendMessage(hwnd, NativeMethods.WM_CLOSE, 0, new IntPtr(0));
               }
            }

            lock (Mutex)
            {
               // Be optimistic.
               this.dialogBoxCloseResult = true;

               for (int i = 0; i < dialogBoxCloseResults.Length; i++)
               {
                  if (!dialogBoxCloseResults[i])
                  {
                     this.dialogBoxCloseResult = false;
                     break;
                  }
               }

               this.threadDone.Set();
            }
         }
      }

      /// <summary>
      /// Finds a messagebox string on a messagebox.
      /// </summary>
      /// <param name="hwnd">The windows handle of the dialog</param>
      /// <param name="unmanagedMemoryLocation">A pointer to the memorylocation the string will be written to</param>
      /// <returns>True if found.</returns>
      private static bool FindMessageBoxString(IntPtr hwnd, IntPtr unmanagedMemoryLocation)
      {
         StringBuilder sb = new StringBuilder(512);
         NativeMethods.GetClassName(hwnd, sb, sb.Capacity);

         if (sb.ToString().ToLower().Contains("static"))
         {
            StringBuilder windowText = new StringBuilder(2048);
            NativeMethods.GetWindowText(hwnd, windowText, windowText.Capacity);

            if (windowText.Length > 0)
            {
               IntPtr stringAsPtr = IntPtr.Zero;
               try
               {
                  stringAsPtr = Marshal.StringToHGlobalAnsi(windowText.ToString());
                  char[] stringAsArray = windowText.ToString().ToCharArray();

                  // Since unicode characters are copied check if we are out of the allocated length.
                  // If not add the end terminating zero.
                  if ((2 * stringAsArray.Length) + 1 < 2048)
                  {
                     Marshal.Copy(stringAsArray, 0, unmanagedMemoryLocation, stringAsArray.Length);
                     Marshal.WriteInt32(unmanagedMemoryLocation, 2 * stringAsArray.Length, 0);
                  }
               }
               finally
               {
                  if (stringAsPtr != IntPtr.Zero)
                  {
                     Marshal.FreeHGlobal(stringAsPtr);
                  }
               }
               return false;
            }
         }

         return true;
      }

      #endregion
   }
}
